"""数据加载模块，生成合成输入数据"""

from __future__ import annotations

from typing import Dict, Any

import numpy as np


def make_synthetic_inputs(model_spec: Dict[str, Any], batch_size: int=1, rng: np.random.Generator=np.random.default_rng()) -> Dict[str, np.ndarray]:
    """
    根据模型规格生成合成输入数据
    
    Args:
        model_spec: 模型规格字典，包含 inputs 字段
        batch_size: 批处理大小，默认1
        rng: NumPy 随机数生成器
        
    Returns:
        输入张量字典 {input_name: np.ndarray}
    """
    inputs = {}
    
    # 支持字典格式
    input_specs = model_spec.get("inputs", {})
    
    if isinstance(input_specs, dict):
        # 字典格式：{"input_name": {"shape": [...], "dtype": "float32", ...}}
        for name, spec in input_specs.items():
            shape = spec["shape"]
            # 如果是可批处理的，替换第一个维度为 batch_size
            if spec.get("batched", True):
                shape = [batch_size] + list(shape[1:])
            else:
                shape = list(shape)
            
            dtype = np.dtype(spec.get("dtype", "float32"))
            
            # 根据数据类型生成合适的随机数据
            if dtype in [np.float32, np.float64]:
                # 浮点数：标准正态分布，然后归一化到 [0, 1]
                data = rng.standard_normal(size=shape).astype(dtype)
                data = (data - data.min()) / (data.max() - data.min() + 1e-8)
            elif dtype in [np.int32, np.int64]:
                # 整数：范围 [0, 1000)，适合 token ids
                data = rng.integers(0, 1000, size=shape, dtype=dtype)
            else:
                # 其他类型：默认随机正态分布
                data = rng.standard_normal(size=shape).astype(dtype)
            
            inputs[name] = data
            
    return inputs


